from __future__ import with_statement

'''
http://www.faqs.org/rfcs/rfc1928.html
'''

import os
import sys
import struct
import random
import socket

from xml.dom.minidom import parse as parseXml, Node as xmlNode
from tempfile import gettempdir
from optparse import OptionParser
from ConfigParser import ParsingError, RawConfigParser as ConfigParser

from twisted.python import log
from twisted.python.runtime import platformType

__package__ = 'pysocks5'
__version__ = '0.2.9'
__authors__ = 'p0is0n (p0is0n@foo.kz)'

parser = OptionParser()

parser.add_option \
(
	'',
	'--pid-file',
	dest = 'pidFile',
	help = 'pid file (default "%default")',
	metavar = 'FILE',
	default = os.path.join(gettempdir(), '%s.pid' % __package__)
)

parser.add_option \
(
	'',
	'--log-file',
	dest = 'logFile',
	help = 'log file (default "%default")',
	metavar = 'FILE',
	default = None#os.path.join(gettempdir(), '%s.log' % __package__)
)

parser.add_option \
(
	'',
	'--config-file',
	dest = 'configFile',
	help = 'config file (default "%default")',
	metavar = 'FILE',
	default = os.path.join(os.path.dirname(__file__), 'config.conf')
)

parser.add_option \
(
	'',
	'--debug',
	dest = 'debug',
	help = 'debug out',
	default = 0,
	action = 'store_const',
	const = 1
)

parser.add_option \
(
	'',
	'--dump',
	dest = 'dump',
	help = 'dump out',
	default = 0,
	action = 'store_const',
	const = 1
)

parser.add_option \
(
	'',
	'--log-level',
	dest = 'logLevel',
	help = 'log level (0 - none, 1 - errors, 2 - warnings,  3 - notices)',
	default = 2,
	type = 'int'
)

parser.add_option \
(
	'',
	'--daemonize',
	dest = 'daemonize',
	help = 'run as daemon',
	default = 0,
	action = 'store_const',
	const = 1
)

options, args = parser.parse_args()

parser.destroy()

(
SERVER,
CLIENT
) = range(2)

def daemonize():
	if platformType == 'posix':
		try:
			if os.fork():
				os._exit(0)
		except OSError, e:
			sys.exit('fork #1 failed %d (%s)\r\n' % (e.errno, e.strerror))
		else:
			os.chdir('/')
			os.umask(0)
			os.setsid()
		try:
			if os.fork():
				os._exit(0)
		except OSError, e:
			sys.exit('fork #2 failed %d (%s)\r\n' % (e.errno, e.strerror))
		else:
			null = os.open('/dev/null', os.O_RDWR)
			for i in range(3):
				try:
					os.dup2(null, i)
				except OSError, e:
					if e.errno != errno.EBADF:
						raise
			os.close(null)
	else:
		sys.exit('fork not allowed, stopping...\r\n')

def checkPid(pidfile):
	if not pidfile:
		return
	if os.path.exists(pidfile):
		try:
			pid = int(open(pidfile).read())
		except ValueError:
			sys.exit('Pidfile %s contains non-numeric value\r\n' % pidfile)
		try:
			if platformType == 'posix':
				os.kill(pid, 0)
		except OSError, why:
			if why[0] == errno.ESRCH:
				sys.stderr.write('Removing stale pidfile %s\r\n' % \
					pidfile)
				os.remove(pidfile)
			else:
				sys.exit('Can\'t check status of PID %s from pidfile %s: %s\r\n' % \
					(pid, pidfile, why[1]))
		else:
			sys.exit('Already running PID %s from pidfile %s\r\n' % \
				(pid, pidfile))

def createPidFile(pidfile):
	if not pidfile:
		return
	with open(pidfile, 'w') as fp:
		fp.write('%s' % os.getpid())

def removePidFile(pidfile):
	if not pidfile:
		return
	if os.path.exists(pidfile):
		os.remove(pidfile)

checkPid(options.pidFile)

if options.debug:
	options.logLevel = 3
if options.logFile and not options.logLevel:
	options.logLevel = 2

options.logLevel > 0 \
	and log.startLogging(open(options.logFile, 'ab') \
		if not options.debug and options.logFile else sys.stdout)

if options.daemonize: daemonize()

try:
	from twisted.internet import kqreactor
	kqreactor.install()
except:
	try:
		from twisted.internet import epollreactor
		epollreactor.install()
	except:
		try:
			from twisted.internet import pollreactor
			pollreactor.install()
		except:
			bestReactor = 'selectreactor'
		else:
			bestReactor = 'pollreactor'
	else:
		bestReactor = 'epollreactor'
else:
	bestReactor = 'kqreactor'

from twisted.internet import protocol, reactor
from twisted.protocols import policies

def dumpPacket(data, rows = 16):
	out = []
	lines = len(data) / rows
	if lines * rows != len(data):
		lines += 1
	for i in xrange(lines):
		d = tuple(data[rows * i:rows * i + rows])
		hex = map(lambda x: '%02X' % ord(x), d)
		text = map(lambda x: (len(repr(x)) > 3 and '.') or x, d)
		out.append(' '.join(hex) + ' ' * 3 * (rows - len(d)) + ' ' + ''.join(text))
	return '\n'.join(out)

class ItemProxyBase(object):

	def __init__(self, host, port):
		self._host = host
		self._port = port
		self._fail = 0
		self._success = 0

	def getType(self):
		"""Return proxy type"""

	def getHost(self):
		return self._host

	def getPort(self):
		return self._port

	def countFails(self):
		return self._fail

	def countSuccess(self):
		return self._success

	def __str__(self):
		return '<Proxy %s:%d>' % \
			(self._host, self._port)

class configFactory(object):

	def __init__(self, file = None):
		if file is not None:
			parser = ConfigParser()
			with open(file, 'rb') as fp:
				parser.readfp(fp)
				for section in parser.sections():
						for key, value in parser.items(section):
							key, value = key.lower(), value.decode('utf8')
							if section == 'general':
								if key in ( \
										'timeout',
										'listenport',
										'connectionslimit',
										):
									value and setattr(self, key, parser.getint(section, key))
								elif key in ( \
										'socksauth',
										):
									setattr(self, key, parser.getboolean(section, key))
								elif key in ( \
										'usersfile',
										'listeninterface',
										):
									value and setattr(self, key, parser.get(section, key))
								else:
									raise ParsingError, 'unknown key "%s" in section "%s"' % \
										(key, section)
							else:
								raise TypeError, 'unknown section "%s"' % \
									(section)

	def getTimeOut(self):
		try:
			return self.timeout
		except:
			return 60

	def getConnectionsLimit(self):
		try:
			return self.connectionslimit
		except:
			return 0

	def getListenPort(self):
		try:
			return self.listenport
		except:
			return 0

	def getListenInterface(self):
		try:
			return str(self.listeninterface)
		except:
			return ''

	def getSocksAuth(self):
		try:
			return self.socksauth
		except:
			return False

	def getUsersFile(self):
		try:
			return str(self.usersfile)
		except:
			return None

	def getAllowInPeers(self):
		pass

	def getAllowOutPeers(self):
		pass

if options.configFile:
	config = configFactory(options.configFile)
else:
	config = configFactory()

class UsersSocks(object):

	users = []

	def __init__(self, file):
		if file:
			xmlUsers = parseXml \
			(
				file
			)
			for element in xmlUsers.getElementsByTagName('user'):
				login, password = None, None
				for node in element.childNodes:
					if node.nodeType == xmlNode.ELEMENT_NODE:
						if node.tagName == 'login':
							login = str(node.lastChild.data)
						elif node.tagName == 'password':
							password = str(node.lastChild.data)
						else:
							raise ValueError('unknown tag "%s" in "%s"' % \
								(node.tagName, options.usersFile))
				login and self.addUser(login, password)
			xmlUsers.unlink()

	def addUser(self, login, password):
		self.users.append((login, password))

	def check(self, login, password):
		code = 0
		try:
			if login:
				for k, v in self.users:
					if k == login:
						if v and password and v == password:
							code = 1
						elif not v and not password:
							code = 1
					if code:
						break
		finally:
			return code

class OutgoingProtocol(protocol.Protocol):

	def __init__(self, father):
		self._father = father
		self._buffer = ''

	def connectionMade(self):
		options.logLevel >= 3 and log.msg('Connection made for server %s' % \
			str(self.transport.getPeer()))
		peer = self.transport.getPeer()
		self._father.makeReply(0x00, port = peer.port, server = peer.host)
		self._father.setOutgoing(self)
		self._father.setState(self._father.STATE_RECEIVE)
		self._father.setTimeout(3600)

	def connectionLost(self, reason):
		options.logLevel >= 3 and log.msg('Connection lost for server %s reason %s' % \
			(str(self.transport.getPeer()), reason or 'unknown'))
		if self._father:
			self._father.setOutgoing(None)
			self._father.transport.loseConnection()

	def dataReceived(self, data):
		options.dump and self.dumpData('server %s [%s] received' % \
			(self.transport.getPeer().host, self), data)
		self._father.write(data)

	def write(self, data):
		options.dump and self.dumpData('server %s [%s] write' % \
			(self.transport.getPeer().host, self), data)
		self.transport.write(data)

	def dumpData(self, *args):
		self._father.dumpData(*args)

class Protocol(policies.TimeoutMixin, protocol.Protocol):

	VERSION = 0x05 # socks version

	STATE_IGNORED = 0x00
	STATE_METHODS = 0x01
	STATE_AUTHREQ = 0x02
	STATE_REQUEST = 0x03
	STATE_RECEIVE = 0x04

	DUMP_I = 0x01
	DUMP_O = 0x02

	AUTH_NONE = 0x00
	AUTH_USPW = 0x02
	AUTH_NSUP = 0xFF

	CODE_TCPC = 0x01
	CODE_TCPB = 0x02
	CODE_UPDA = 0x03

	TYPE_IPv4 = 0x01
	TYPE_IPv6 = 0x04
	TYPE_DOMN = 0x03

	outgoingClass = OutgoingProtocol

	def __init__(self):
		self._authTypes = []
		if not config.getSocksAuth():
			self._authTypes.append(self.AUTH_NONE)
		self._authTypes.append(self.AUTH_USPW)
		self._sessionState = None
		self._selectedMethod = None
		self._buffer = ''
		self._outgoing = None

	def dataReceived(self, data):
		self.resetTimeout()
		self._buffer = self._buffer + data
		if not self.isState(self.STATE_RECEIVE):
			options.dump and self.dumpData('client %s [%s] received' % \
				(self.transport.getPeer().host, self), data)
		if self._sessionState is None:
			self.setState(self.STATE_METHODS)
		if self.isState(self.STATE_IGNORED):
			return
		if self.isState(self.STATE_METHODS):
			if len(self._buffer) < 3:
				return
			try:
				version, count = struct.unpack('!2B', self.fromBuffer(2))
			except struct.error, e:
				options.logLevel >= 1 and log.msg('Error: struct error %s' % repr(e))
				return self.transport.loseConnection()
			if version != self.VERSION:
				options.logLevel >= 2 and log.msg('Warning: Wrong version from client %s' % \
					str(self.transport.getPeer()))
				return self.transport.loseConnection()
			try:
				methods = struct.unpack('!%db' % count, self.fromBuffer(count))
			except struct.error, e:
				options.logLevel >= 1 and log.msg('Error: struct error %s' % repr(e))
				return self.transport.loseConnection()
			for method in methods:
				if method in self._authTypes:
					self._selectedMethod = method
					break
				else:
					self._selectedMethod = None
			if self._selectedMethod is None:
				options.logLevel >= 3 and log.msg('NO ACCEPTABLE METHODS [%r] from %s' % \
					(methods, str(self.transport.getPeer())))
				responseCode = self.AUTH_NSUP # NO ACCEPTABLE METHODS
			else:
				responseCode = self._selectedMethod
			if responseCode == self.AUTH_NONE:
				self.setState(self.STATE_REQUEST)
			elif responseCode == self.AUTH_USPW:
				self.setState(self.STATE_AUTHREQ)
			self.write(struct.pack('!2B', self.VERSION, responseCode))
			if self._selectedMethod is None:
				return self.transport.loseConnection()
		elif self.isState(self.STATE_AUTHREQ):
			try:
				version, loginLength = struct.unpack('!2B', self.fromBuffer(2))
			except struct.error, e:
				options.logLevel >= 1 and log.msg('Error: struct error %s' % repr(e))
				return self.transport.loseConnection()
			if version != 0x01:
				options.logLevel >= 2 and log.msg('Warning: Wrong version for auth from client %s' % \
					str(self.transport.getPeer()))
				return self.transport.loseConnection()
			try:
				login, passwordLength = struct.unpack('!%dsB' % loginLength, self.fromBuffer(1 + loginLength))
			except struct.error, e:
				options.logLevel >= 1 and log.msg('Error: struct error %s' % repr(e))
				return self.transport.loseConnection()
			if not usersSocks:
				responseCode = 1
			else:
				responseCode = not usersSocks.check \
				(
					login, self.fromBuffer(passwordLength)
				)
			self.write(struct.pack('!2B', 0x01, responseCode))
			if not responseCode:
				self.setState(self.STATE_REQUEST)
			else:
				options.logLevel >= 3 and log.msg('Wrong login from %s' % \
					str(self.transport.getPeer()))
				self.transport.loseConnection()
		elif self.isState(self.STATE_REQUEST):
			try:
				version, code, byte, self._type = struct.unpack('!4B', self.fromBuffer(4))
			except struct.error, e:
				options.logLevel >= 1 and log.msg('Error: struct error %s' % repr(e))
				return self.transport.loseConnection()
			if version != self.VERSION:
				options.logLevel >= 2 and log.msg('Warning: Wrong version from client %s' % \
					str(self.transport.getPeer()))
				return self.transport.loseConnection()
			if self._type == self.TYPE_IPv4:
				server = socket.inet_ntoa(self.fromBuffer(4))
			elif self._type == self.TYPE_DOMN:
				lenght = struct.unpack('!B', self.fromBuffer(1))[0]
				server = self.fromBuffer(lenght)
				if not server:
					options.logLevel >= 2 and log.msg('Warning: missing domain for client %s' % \
						self.transport.getPeer().host)
					return self.makeReply(code = 0x01, end = True)
			else:
				return self.makeReply(code = 0x08, end = True) # Address type not supported
			port = struct.unpack('!H', self.fromBuffer(2))[0]
			if code == self.CODE_TCPC:
				d = self.connectClass(server, port, self.outgoingClass, self)
				d.addErrback(lambda result, self=self: self.makeReply(code=0x04, end=True))
			else:
				return self.makeReply(code=0x01, end=True)
			self.setState(self.STATE_IGNORED)
		elif self.isState(self.STATE_RECEIVE):
			options.dump and self.dumpData('client %s [%s] received' % \
				(self.transport.getPeer().host, self), self._buffer)
			self._outgoing.write(self._buffer)
			self._buffer = ''
		else:
			try:
				return self.transport.loseConnection()
			finally:
				options.logLevel >= 1 \
					and log.msg('Error: WTF?! where state 0x%02X' % self._sessionState)
		self.setTimeout(self.timeOut)

	def isState(self, state):
		return self._sessionState == state

	def setState(self, state):
		self._sessionState = state

	def fromBuffer(self, bytes=None):
		if bytes is None:
			bytes = len(self._buffer)
		data, self._buffer = \
			self._buffer[:bytes], self._buffer[bytes:]
		return data

	def makeReply(self, code, port=0, server=None, end=False):
		if self._type == self.TYPE_DOMN:
			if server is None:
				server = ''
			self.write(struct.pack('!5B%dsH' % len(server), self.VERSION, code, 0x00, \
				self._type, len(server), server, port))
		elif self._type == self.TYPE_IPv4:
			if server is None:
				server = '0.0.0.0'
			self.write(struct.pack('!4B', self.VERSION, code, 0x00, self._type) + \
				socket.inet_aton(server) + struct.pack('!H', port))
		end and self.transport.loseConnection()

	def write(self, data):
		options.dump and self.dumpData('client %s [%s] write' % \
			(self.transport.getPeer().host, self), data)
		self.transport.write(data)

	def connectClass(self, host, port, klass, *args):
		 return protocol.ClientCreator \
		 (
		 	reactor, klass, *args
		 ).connectTCP(host, port)

	def timeoutConnection(self):
		options.logLevel >= 3 and log.msg('Timing out client %s' % \
			str(self.transport.getPeer()))
		policies.TimeoutMixin.timeoutConnection(self)

	def connectionMade(self):
		self.factory.registerProtocol(self)
		options.logLevel >= 3 and log.msg('Connection made for client %s' % \
			str(self.transport.getPeer()))
		self.setTimeout(self.timeOut)

	def connectionLost(self, reason):
		self.factory.unregisterProtocol(self)
		options.logLevel >= 3 and log.msg('Connection lost for client %s reason %s' % \
			(str(self.transport.getPeer()), reason or 'unknown'))
		self.setTimeout(None)
		if self._outgoing:
			self._outgoing.transport.loseConnection()
		self._outgoing = None

	def setOutgoing(self, outgoing):
		self._outgoing = outgoing

	def dumpData(self, type, data):
		log.msg('Dump %s' % type)
		log.msg('\n\n%s\n' % dumpPacket(data))

class Factory(policies.LimitTotalConnectionsFactory):

	protocol = Protocol

	connectionLimit = config.getConnectionsLimit()

	def registerProtocol(self, p):
		p.timeOut = config.getTimeOut()

reactor.addSystemEventTrigger \
(
	'before', 'shutdown', removePidFile, options.pidFile
)

reactor.listenTCP \
(
	interface = config.getListenInterface(), \
		port = config.getListenPort(), factory = Factory()
)

reactor.callWhenRunning \
(
	createPidFile, options.pidFile
)

reactor.callWhenRunning \
(
	log.msg, '%s %s author %s log level %d reactor %s' % \
		(__package__, __version__, __authors__, options.logLevel, bestReactor)
)

reactor.callWhenRunning \
(
	log.msg, 'config file %s' % \
		(options.configFile)
)

if config.getUsersFile():
	reactor.callWhenRunning \
	(
		log.msg, 'users file %s' % \
			(config.getUsersFile())
	)
elif config.getSocksAuth():
	reactor.callWhenRunning \
	(
		log.msg, 'Warning: you not select users file, please edit config %s' % \
			(options.configFile)
	)

if config.getSocksAuth():
	usersSocks = UsersSocks \
	(
		config.getUsersFile()
	)
else:
	usersSocks = None

reactor.callWhenRunning \
(
	log.msg, 'Ok, start'
)

reactor.addSystemEventTrigger \
(
	'before', 'shutdown', log.msg, 'Ok, stops'
)

reactor.run()
